package protocol

import (
	"context"
	"errors"
	"fmt"
	"math"
	"sync"
	"time"

	"gitee.com/doraemon1996/bunny/structure"
	"gitee.com/doraemon1996/bunny/timer"
	"gitee.com/doraemon1996/bunny/uuid"
	"google.golang.org/protobuf/proto"
)

var (
	// 消息包管理器已关闭
	ErrPacketMngClosed = errors.New("packet manager closed")
)

// 消息包管理器对象
type PacketMng interface {
	// 初始化消息包管理器
	//	maxLength: 消息包最大长度
	//	expireTime: 组包过期时间
	Init(ctx context.Context, maxLength int, expireTime time.Duration)

	// 启动管理器
	Start()

	// 停止管理器
	Stop()

	// 消息包编码
	//  msg: 消息结构
	//  list: 编码且根据消息包最大长度分段后的消息数据
	Pack(msg *Message) (list [][]byte, err error)

	// 消息包解码
	//	data: 消息数据
	//	msg: 解码后的消息结构
	//	wait: 消息包是否需要等待组包 true表示需要
	Unpack(data []byte, msg *Message) (wait bool, err error)
}

// 消息包管理器
type packetMng struct {
	structure.Control

	packetMaxLength    int           // 消息包最大长度
	assembleExpireTime time.Duration // 组包过期时间
	crontabID          uint64

	assemblerPool *structure.CachePool[*assembler] // 组包器缓存池
	packetPool    *structure.CachePool[*Packet]    // 消息包缓存池

	muRecive sync.Mutex            // 收包等待表操作锁
	receive  map[uint64]*assembler // 收包等待表
}

// 创建消息包管理器
func NewPacketMng() PacketMng {
	return new(packetMng)
}

func (pm *packetMng) Init(ctx context.Context, maxLength int, expireTime time.Duration) {
	pm.Control.Init(ctx, func() {
		pm.packetMaxLength = maxLength
		if expireTime < time.Second {
			expireTime = time.Second
		}
		pm.assembleExpireTime = expireTime

		pm.assemblerPool = structure.NewCachePool[*assembler](func() any {
			return new(assembler)
		})

		pm.packetPool = structure.NewCachePool[*Packet](func() any {
			return new(Packet)
		})

		pm.receive = make(map[uint64]*assembler)
	})
}

func (pm *packetMng) Start() {
	pm.Control.Start(func() {
		pm.crontabID = timer.Interval(time.Second, func() {
			pm.expire()
		})
	})
}

func (pm *packetMng) Stop() {
	pm.Control.Stop(func() {
		timer.Remove(pm.crontabID)

		pm.Wait()

		pm.muRecive.Lock()
		for _, assembler := range pm.receive {
			pm.assemblerPool.Put(assembler)
		}
		pm.receive = nil
		pm.muRecive.Unlock()
	})
}

// 收包等待列表过期操作
func (pm *packetMng) expire() {
	if pm.Closed() {
		return
	}

	pm.Add(1)
	defer pm.Done()

	now := time.Now()

	pm.muRecive.Lock()
	for packetID, assembler := range pm.receive {
		if assembler.Expired(now, pm.assembleExpireTime) {
			delete(pm.receive, packetID)
			pm.assemblerPool.Put(assembler)
		}
	}
	pm.muRecive.Unlock()
}

// 组包
func (pm *packetMng) assemble(pkt *Packet) (list []*Packet, dataLen int) {
	pm.muRecive.Lock()
	assembler, exist := pm.receive[pkt.PacketId]
	if !exist {
		assembler = pm.assemblerPool.Get()
		assembler.Init()
		pm.receive[pkt.PacketId] = assembler
	}

	complete := assembler.Add(pkt)
	if complete {
		delete(pm.receive, pkt.PacketId)
	}
	pm.muRecive.Unlock()

	if complete {
		list, dataLen = assembler.Assemble()
		pm.assemblerPool.Put(assembler)
	}

	return
}

// 生成消息包ID
func (pm *packetMng) getPacketID() uint64 {
	return uuid.GenUint64()
}

// 封包
func (pm *packetMng) Pack(message *Message) (list [][]byte, err error) {
	if pm.Closed() {
		err = ErrPacketMngClosed
		return
	}

	pm.Add(1)
	defer pm.Done()

	msgData, msgErr := proto.Marshal(message)
	if msgErr != nil {
		err = fmt.Errorf("proto.Marshal message err: %v", msgErr)
		return
	}

	packetID := pm.getPacketID()

	// 小于等于单包最大限制
	if len(msgData) <= pm.packetMaxLength {
		packet := pm.packetPool.Get()
		packet.PacketId = packetID
		packet.ChunkSeq = 0
		packet.ChunkCount = 1
		packet.Chunk = msgData
		pktData, pktErr := proto.Marshal(packet)
		pm.packetPool.Put(packet)
		if pktErr != nil {
			err = fmt.Errorf("proto.Marshal packet err: %v", pktErr)
		} else {
			list = append(list, pktData)
		}
		return
	}

	// 分块数量
	section := int(math.Ceil(float64(len(msgData)) / float64(pm.packetMaxLength)))
	bodyLen := int(math.Ceil(float64(len(msgData)) / float64(section)))
	buffer := structure.NewBuffer(msgData)
	for i := 1; i <= section; i++ {
		packet := pm.packetPool.Get()
		packet.PacketId = packetID
		packet.ChunkSeq = uint64(i)
		packet.ChunkCount = uint64(section)
		packet.Chunk = buffer.Next(bodyLen)
		pktData, pktErr := proto.Marshal(packet)
		pm.packetPool.Put(packet)
		if pktErr != nil {
			err = fmt.Errorf("proto.Marshal packet err: %v", pktErr)
			return
		} else {
			list = append(list, pktData)
		}
	}
	return
}

// 解包
// wait 表示消息被分块了, 等待所有块都被接收时才会解包
func (pm *packetMng) Unpack(data []byte, msg *Message) (wait bool, err error) {
	if pm.Closed() {
		err = ErrPacketMngClosed
		return
	}

	pm.Add(1)
	defer pm.Done()

	packet := pm.packetPool.Get()
	if dataErr := proto.Unmarshal(data, packet); dataErr != nil {
		err = fmt.Errorf("proto.Unmarshal packet err: %v", dataErr)
		pm.packetPool.Put(packet)
		return
	}

	var msgData []byte
	if packet.ChunkCount > 1 {
		list, dataLen := pm.assemble(packet)
		wait = len(list) == 0
		if !wait {
			buffer := structure.NewBuffer(make([]byte, 0, dataLen))
			for _, packet := range list {
				buffer.Write(packet.Chunk)
				pm.packetPool.Put(packet)
			}
			msgData = buffer.Bytes()
		}
	} else {
		msgData = packet.Chunk
		pm.packetPool.Put(packet)
	}

	// 需要粘包
	if wait {
		return
	}

	if pktErr := proto.Unmarshal(msgData, msg); pktErr != nil {
		err = fmt.Errorf("proto.Unmarshal message err: %v", pktErr)
	}

	return
}
